# 2025年2月20日10:17:36
# 将json格式的标注转为yolo格式

import os
import xml.etree.ElementTree as ET
from PIL import Image

# 定义类别映射
class_mapping = {
    'helmet': 0,
    'without_helmet': 1,
    'two_wheeler': 2
}

# 定义文件夹路径
base_folder = r'D:\上过的课\2024-2025-1-大四上\毕设\datasets\TWHD'
annotations_folder = os.path.join(base_folder, 'Annotations')
images_folder = os.path.join(base_folder, 'JPEGImages')
labels_folder = os.path.join(base_folder, 'labels')

# 创建labels文件夹
if not os.path.exists(labels_folder):
    os.makedirs(labels_folder)

# 遍历JPEGImages文件夹中的所有图片
for image_name in os.listdir(images_folder):
    if image_name.endswith(('.jpg', '.jpeg', '.png')):
        # 获取图片的文件名（不包含扩展名）
        image_base_name = os.path.splitext(image_name)[0]
        # 构建对应的标注文件路径
        xml_file = os.path.join(annotations_folder, f'{image_base_name}.xml')
        # 构建对应的YOLO标注文件路径
        yolo_file = os.path.join(labels_folder, f'{image_base_name}.txt')

        # 打开图片获取其宽度和高度
        image_path = os.path.join(images_folder, image_name)
        with Image.open(image_path) as img:
            img_width, img_height = img.size

        # 解析XML文件
        if os.path.exists(xml_file):
            tree = ET.parse(xml_file)
            root = tree.getroot()

            # 打开YOLO标注文件以写入标注信息
            with open(yolo_file, 'w') as f:
                for obj in root.findall('object'):
                    # 获取类别名称
                    class_name = obj.find('name').text
                    if class_name in class_mapping:
                        class_id = class_mapping[class_name]
                        # 获取边界框信息
                        bbox = obj.find('bndbox')
                        xmin = int(bbox.find('xmin').text)
                        ymin = int(bbox.find('ymin').text)
                        xmax = int(bbox.find('xmax').text)
                        ymax = int(bbox.find('ymax').text)

                        # 计算YOLO格式的边界框信息
                        x_center = (xmin + xmax) / (2 * img_width)
                        y_center = (ymin + ymax) / (2 * img_height)
                        width = (xmax - xmin) / img_width
                        height = (ymax - ymin) / img_height

                        # 写入YOLO格式的标注信息
                        f.write(f'{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n')

print("转换完成！")